#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
统计ZIP文件中图片文件数量的工具
支持统计.jpg和.png格式的图片文件
"""

import argparse
import zipfile
import os
from pathlib import Path


def count_images_in_zip(zip_path):
    """
    统计ZIP文件中图片文件的数量
    
    Args:
        zip_path (str): ZIP文件路径
        
    Returns:
        int: 图片文件数量
    """
    image_count = 0
    image_extensions = {'.jpg', '.jpeg', '.png'}
    
    try:
        # 检查文件是否存在
        if not os.path.exists(zip_path):
            raise FileNotFoundError(f"文件不存在: {zip_path}")
        
        # 检查是否为ZIP文件
        if not zipfile.is_zipfile(zip_path):
            raise ValueError(f"文件不是有效的ZIP格式: {zip_path}")
        
        # 打开ZIP文件并统计图片
        with zipfile.ZipFile(zip_path, 'r') as zip_file:
            for file_info in zip_file.infolist():
                # 忽略目录
                if not file_info.is_dir():
                    # 获取文件扩展名并转换为小写
                    file_extension = Path(file_info.filename).suffix.lower()
                    if file_extension in image_extensions:
                        image_count += 1
                        
        return image_count
    
    except FileNotFoundError as e:
        print(f"错误: {e}")
        return -1
    except zipfile.BadZipFile:
        print(f"错误: 文件不是有效的ZIP文件: {zip_path}")
        return -1
    except Exception as e:
        print(f"处理ZIP文件时出错: {e}")
        return -1


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='统计ZIP文件中图片文件数量')
    parser.add_argument('zip_file', help='ZIP文件路径')
    
    args = parser.parse_args()
    
    print(f"正在处理ZIP文件: {args.zip_file}")
    count = count_images_in_zip(args.zip_file)
    
    if count >= 0:
        print(f"ZIP文件中包含 {count} 个图片文件 (.jpg, .jpeg, .png)")
    else:
        print("处理失败")


if __name__ == "__main__":
    main()